package com.example.scribe.backend;

import java.util.ArrayList;
import java.util.List;

import com.example.scribe.SyntaxNode;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackManipulation.Compound;
import com.example.scribe.backend.handlers.*;

/**
 * Basic implementation of StatementEmitter using Strategy pattern with handlers.
 * Each statement type is handled by a dedicated handler registered in the registry.
 */
public final class BasicStatementEmitter implements StatementEmitter {
    private final NodeHandlerRegistry registry;

    BasicStatementEmitter(ExpressionEmitter exprEmitter) {
        this.registry = new NodeHandlerRegistry();
        initializeHandlers(exprEmitter);
    }

    /**
     * Get the registry for accessing handlers.
     * Used by WhileStatementHandler to process nested statements.
     */
    public NodeHandlerRegistry getRegistry() {
        return registry;
    }

    /**
     * Initialize and register all statement handlers.
     */
    private void initializeHandlers(ExpressionEmitter exprEmitter) {
        registry.registerStatementHandler("IfStatement", new IfStatementHandler());
        registry.registerStatementHandler("return_stmt", new ReturnStatementHandler());
        registry.registerStatementHandler("VariableDecl", new VariableDeclStatementHandler());
        registry.registerStatementHandler("WhileStatement", new WhileStatementHandler());
        registry.registerStatementHandler("assignment", new com.example.scribe.backend.handlers.AssignmentStatementHandler());
    }

    @Override
    public StackManipulation emitBody(String functionName, SyntaxNode bodyNode, String generatedClassName, LocalVariableAllocator locals, ExpressionEmitter exprEmitter) {
        if (bodyNode == null) return defaultZeroReturn();


        // Special case: handle simple pattern "let x = expr; return x;"
        String varName = null;
        SyntaxNode varInitExpr = null;
        List<SyntaxNode> stmts = new ArrayList<>();
        for (SyntaxNode s : bodyNode.children()) {
            if ("statement".equals(s.type())) {
                stmts.add(s);
            }
        }
        if (stmts.size() >= 2) {
            SyntaxNode first = stmts.get(0);
            for (SyntaxNode fc : first.children()) {
                if ("VariableDecl".equals(fc.type())) {
                    for (SyntaxNode vc : fc.children()) {
                        if (ASTUtils.isIdentifier(vc)) {
                            varName = vc.text();
                        }
                        if ("expression".equals(vc.type())) {
                            varInitExpr = vc;
                        }
                    }
                }
            }
            SyntaxNode second = stmts.get(1);
            boolean secondIsReturnOfVar = false;
            for (SyntaxNode sc : second.children()) {
                if ("return_stmt".equals(sc.type())) {
                    // Check if return expression contains the variable name
                    SyntaxNode returnExpr = ASTUtils.findFirst(sc, "expression");
                    if (returnExpr != null) {
                        // Use findIdentifierOrNumber to check if it matches varName
                        // BUT: only match if return expression is JUST the variable (not a complex expression)
                        // Check if returnExpr directly contains only the variable identifier
                        boolean isSimpleVarReturn = false;
                        List<SyntaxNode> returnNodes = new ArrayList<>();
                        ASTUtils.collectAllNodes(returnExpr, returnNodes);
                        // Count how many identifier nodes match varName
                        int matchingIdentifiers = 0;
                        for (SyntaxNode rn : returnNodes) {
                            if (ASTUtils.isIdentifier(rn) && varName != null && varName.equals(rn.text())) {
                                matchingIdentifiers++;
                            }
                        }
                        // Only match if there's exactly one matching identifier AND no structured operations
                        // (shift, unary, bitwise, arithmetic, etc.)
                        boolean hasStructuredOps = false;
                        for (SyntaxNode rn : returnNodes) {
                            String rnType = rn.type();
                            if ("shift".equals(rnType) || "unary".equals(rnType) ||
                                "bitwise_or".equals(rnType) || "bitwise_xor".equals(rnType) ||
                                "bitwise_and".equals(rnType) || "addition".equals(rnType) ||
                                "multiplication".equals(rnType) || "comparison".equals(rnType) ||
                                "equality".equals(rnType)) {
                                hasStructuredOps = true;
                                break;
                            }
                        }
                        if (matchingIdentifiers == 1 && !hasStructuredOps) {
                            secondIsReturnOfVar = true;
                        }
                    }
                }
            }
            if (varName != null && varInitExpr != null && secondIsReturnOfVar) {
                int localIndex = locals.allocateLocal(varName);
                StackManipulation init = exprEmitter.emit(varInitExpr, generatedClassName, locals);
                if (init == StackManipulation.Trivial.INSTANCE) {
                    StackManipulation fallback = findIdentifierOrNumber(varInitExpr, generatedClassName, locals);
                    if (fallback != StackManipulation.Trivial.INSTANCE) {
                        init = fallback;
                    } else {
                        init = BytecodeOps.pushInt(0);
                    }
                }
                return new Compound(
                    init,
                    BytecodeOps.storeInt(localIndex),
                    BytecodeOps.loadInt(localIndex),
                    BytecodeOps.ireturn()
                );
            }
        }

        // Collect leading variable declarations
        List<StackManipulation> leadingStores = new ArrayList<>();
        for (SyntaxNode s : bodyNode.children()) {
            if (!"statement".equals(s.type())) continue;
            
            // Try to find VariableDecl at any depth in this statement
            SyntaxNode decl = ASTUtils.findFirst(s, "VariableDecl");
            // Also check direct children
            for (SyntaxNode sc : s.children()) {
                if ("VariableDecl".equals(sc.type())) {
                    decl = sc;
                    break;
                }
            }
            
            if (decl != null) {
                // Use handler to process variable declaration
                StatementNodeHandler handler = registry.getStatementHandler("VariableDecl");
                if (handler != null && handler.canHandle("VariableDecl")) {
                    StackManipulation varDeclEmit = handler.handle(decl, generatedClassName, locals, exprEmitter, this);
                    if (varDeclEmit != StackManipulation.Trivial.INSTANCE) {
                        leadingStores.add(varDeclEmit);
                    }
                }
            }
        }

        // Process statements using handlers
        for (SyntaxNode stmt : bodyNode.children()) {
            if ("statement".equals(stmt.type())) {
                for (SyntaxNode child : stmt.children()) {
                    String childType = child.type();
                    
                    // Handle IfStatement
                    if ("IfStatement".equals(childType)) {
                        StatementNodeHandler handler = registry.getStatementHandler("IfStatement");
                        if (handler != null && handler.canHandle("IfStatement")) {
                            StackManipulation branch = handler.handle(child, generatedClassName, locals, exprEmitter, this);
                            if (branch != StackManipulation.Trivial.INSTANCE) {
                                if (leadingStores.isEmpty()) {
                                    return branch;
                                }
                                List<StackManipulation> seq = new ArrayList<>(leadingStores);
                                seq.add(branch);
                                return new Compound(seq.toArray(new StackManipulation[0]));
                            }
                        }
                        // unknown If shape: skip specialized handling, let generic sequencing fall through
                        continue;
                    }

                    // Handle WhileStatement
                    // IMPORTANT: WhileStatement should NOT cause early return - it should be part of a sequence
                    // The loop itself doesn't return, so we need to continue processing subsequent statements
                    if ("WhileStatement".equals(childType)) {
                        StatementNodeHandler handler = registry.getStatementHandler("WhileStatement");
                        if (handler != null && handler.canHandle("WhileStatement")) {
                            StackManipulation loop = handler.handle(child, generatedClassName, locals, exprEmitter, this);
                            if (loop != StackManipulation.Trivial.INSTANCE) {
                                // Add loop to sequence, but don't return - continue processing subsequent statements
                                leadingStores.add(loop);
                            }
                        }
                        continue;
                    }
                    
                    // Handle assignment (may appear as expression_stmt)
                    if ("assignment".equals(childType)) {
                        StatementNodeHandler handler = registry.getStatementHandler("assignment");
                        if (handler != null && handler.canHandle("assignment")) {
                            StackManipulation assign = handler.handle(child, generatedClassName, locals, exprEmitter, this);
                            if (assign != StackManipulation.Trivial.INSTANCE) {
                                // Add to sequence, but don't return - continue processing
                                leadingStores.add(assign);
                            }
                        }
                        continue;
                    }
                    
                    // Handle expression_stmt that might contain assignment
                    if ("expression_stmt".equals(childType)) {
                        // Check if expression_stmt contains assignment
                        SyntaxNode assignment = ASTUtils.findFirst(child, "assignment");
                        if (assignment != null) {
                            StatementNodeHandler handler = registry.getStatementHandler("assignment");
                            if (handler != null && handler.canHandle("assignment")) {
                                StackManipulation assign = handler.handle(assignment, generatedClassName, locals, exprEmitter, this);
                                if (assign != StackManipulation.Trivial.INSTANCE) {
                                    leadingStores.add(assign);
                                }
                            }
                        }
                        continue;
                    }
                    
                    // Handle return_stmt
                    if ("return_stmt".equals(childType)) {
                        StatementNodeHandler handler = registry.getStatementHandler("return_stmt");
                        if (handler != null && handler.canHandle("return_stmt")) {
                            StackManipulation returnEmit = handler.handle(child, generatedClassName, locals, exprEmitter, this);
                            // Prepend leadingStores if any
                            List<StackManipulation> seq = new ArrayList<>();
                            seq.addAll(leadingStores);
                            seq.add(returnEmit);
                            return new Compound(seq.toArray(new StackManipulation[0]));
                        }
                    }
                }
            }
        }

        // Generic block sequencing: emit variable declarations (stores) preceding the first return, then return expr
        List<SyntaxNode> stmts2 = new ArrayList<>();
        for (SyntaxNode s : bodyNode.children()) {
            if ("statement".equals(s.type())) {
                stmts2.add(s);
            }
        }
        int returnIdx = -1;
        for (int i = 0; i < stmts2.size(); i++) {
            if (ASTUtils.containsType(stmts2.get(i), "return_stmt")) {
                returnIdx = i;
                break;
            }
        }
        if (returnIdx >= 0) {
            List<StackManipulation> seq = new ArrayList<>();
            for (int i = 0; i < returnIdx; i++) {
                SyntaxNode s = stmts2.get(i);
                SyntaxNode decl = ASTUtils.findFirst(s, "VariableDecl");
                if (decl != null) {
                    StatementNodeHandler handler = registry.getStatementHandler("VariableDecl");
                    if (handler != null && handler.canHandle("VariableDecl")) {
                        StackManipulation varDeclEmit = handler.handle(decl, generatedClassName, locals, exprEmitter, this);
                        if (varDeclEmit != StackManipulation.Trivial.INSTANCE) {
                            seq.add(varDeclEmit);
                        }
                    }
                }
            }
            SyntaxNode ret = stmts2.get(returnIdx);
            SyntaxNode rexpr = ASTUtils.findFirst(ret, "expression");
            if (rexpr != null) {
                StackManipulation re = exprEmitter.emit(rexpr, generatedClassName, locals);
                if (re == StackManipulation.Trivial.INSTANCE) {
                    // Try to find identifier or number recursively in expression subtree
                    StackManipulation fallback = findIdentifierOrNumber(rexpr, generatedClassName, locals);
                    if (fallback != StackManipulation.Trivial.INSTANCE) {
                        re = fallback;
                    } else {
                        re = BytecodeOps.pushInt(0);
                    }
                }
                // Ensure re is not Trivial
                if (re == StackManipulation.Trivial.INSTANCE) {
                    re = BytecodeOps.pushInt(0);
                }
                seq.add(re);
                seq.add(BytecodeOps.ireturn());
                return new Compound(seq.toArray(new StackManipulation[0]));
            }
        }

        return defaultZeroReturn();
    }

    private StackManipulation defaultZeroReturn() {
        return new Compound(BytecodeOps.pushInt(0), BytecodeOps.ireturn());
    }

    /**
     * Helper method to find identifier or number in expression subtree.
     */
    private StackManipulation findIdentifierOrNumber(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals) {
        if (node == null) return StackManipulation.Trivial.INSTANCE;
        
        // Check for identifier as local variable
        if (ASTUtils.isIdentifier(node)) {
            String name = node.text();
            if (name != null && locals.hasName(name)) {
                return BytecodeOps.loadInt(locals.getIndexByName(name));
            }
        }
        
        // Check for number
        if (ASTUtils.isNumber(node)) {
            try {
                int v = Integer.parseInt(node.text());
                return BytecodeOps.pushInt(v);
            } catch (NumberFormatException e) {
                // continue
            }
        }
        
        // Also check for identifier in primary nodes
        if ("primary".equals(node.type())) {
            for (SyntaxNode child : node.children()) {
                if (ASTUtils.isIdentifier(child)) {
                    String name = child.text();
                    if (name != null && locals.hasName(name)) {
                        return BytecodeOps.loadInt(locals.getIndexByName(name));
                    }
                }
            }
        }
        
        // Recursively search children
        for (SyntaxNode child : node.children()) {
            StackManipulation found = findIdentifierOrNumber(child, generatedClassName, locals);
            if (found != StackManipulation.Trivial.INSTANCE) return found;
        }
        
        return StackManipulation.Trivial.INSTANCE;
    }
}
